# TODO: Crawling sample for existing active Osv intelligence

import os
import yaml
import git
import shutil
from typing import Dict, List, Optional
from pathlib import Path
from loguru import logger
from datetime import datetime

class GitScraper:
    """
    Git 仓库样本抓取器
    用于从配置的 Git 仓库中抓取与 OSV 情报相关的样本
    """

    def __init__(self, config: Dict):
        """
        初始化 Git 抓取器
        :param config: 配置信息
        """
        self.config = config
        self.temp_dir = Path(config['storage']['temp_dir'])
        self.data_dir = Path(config['storage']['data_dir'])
        self._setup_directories()

    def _setup_directories(self):
        """创建必要的目录"""
        self.temp_dir.mkdir(parents=True, exist_ok=True)
        self.data_dir.mkdir(parents=True, exist_ok=True)

    def scrape_repository(self, source: Dict) -> Optional[Path]:
        """
        抓取单个仓库的样本
        :param source: 仓库配置信息
        :return: 样本存储路径
        """
        repo_name = source['name']
        repo_url = source['url']
        branch = source.get('branch', 'main')
        
        try:
            # 创建临时目录
            temp_repo_dir = self.temp_dir / repo_name
            if temp_repo_dir.exists():
                shutil.rmtree(temp_repo_dir)
            
            # 克隆仓库
            logger.info(f"正在克隆仓库: {repo_url}")
            repo = git.Repo.clone_from(repo_url, temp_repo_dir, branch=branch)
            
            # 切换到指定分支
            if branch != 'main':
                repo.git.checkout(branch)
            
            # 处理样本
            samples_dir = self._process_samples(temp_repo_dir, repo_name)
            
            # 清理临时目录
            shutil.rmtree(temp_repo_dir)
            
            return samples_dir
            
        except Exception as e:
            logger.error(f"处理仓库 {repo_url} 时发生错误: {str(e)}")
            return None

    def _process_samples(self, repo_dir: Path, repo_name: str) -> Path:
        """
        处理仓库中的样本
        :param repo_dir: 仓库目录
        :param repo_name: 仓库名称
        :return: 样本存储路径
        """
        # 创建样本存储目录
        samples_dir = self.data_dir / repo_name
        samples_dir.mkdir(parents=True, exist_ok=True)
        
        # 根据仓库类型处理样本
        if repo_name == "github-malicious-packages":
            self._process_github_malicious_packages(repo_dir, samples_dir)
        elif repo_name == "datadog-malicious-packages":
            self._process_datadog_malicious_packages(repo_dir, samples_dir)
        elif repo_name == "github-advisory-database":
            self._process_github_advisory_database(repo_dir, samples_dir)
        
        return samples_dir

    def _process_github_malicious_packages(self, repo_dir: Path, samples_dir: Path):
        """处理 GitHub 恶意包数据"""
        # 复制 JSON 文件
        for json_file in repo_dir.glob("**/*.json"):
            if "malicious" in json_file.name.lower():
                shutil.copy2(json_file, samples_dir / json_file.name)

    def _process_datadog_malicious_packages(self, repo_dir: Path, samples_dir: Path):
        """处理 DataDog 恶意包数据"""
        # 复制数据集文件
        for data_file in repo_dir.glob("**/*.json"):
            if "dataset" in data_file.name.lower():
                shutil.copy2(data_file, samples_dir / data_file.name)

    def _process_github_advisory_database(self, repo_dir: Path, samples_dir: Path):
        """处理 GitHub 安全公告数据库"""
        # 复制安全公告文件
        for advisory_file in repo_dir.glob("**/*.json"):
            if "GHSA" in advisory_file.name:
                shutil.copy2(advisory_file, samples_dir / advisory_file.name)

    @classmethod
    def scrape_all(cls, config: Dict) -> List[Path]:
        """
        抓取所有配置的仓库样本
        :param config: 配置信息
        :return: 样本存储路径列表
        """
        scraper = cls(config)
        sample_dirs = []
        
        for source in config['osv_sources']:
            try:
                sample_dir = scraper.scrape_repository(source)
                if sample_dir:
                    sample_dirs.append(sample_dir)
                    logger.info(f"成功抓取仓库 {source['name']} 的样本")
            except Exception as e:
                logger.error(f"抓取仓库 {source['name']} 时发生错误: {str(e)}")
                continue
        
        return sample_dirs
